#!/bin/bash

# 检查是否以root用户运行
if [ "$EUID" -ne 0 ]; then
  echo "请以root用户运行此脚本"
  exit
fi

# 配置文件路径
CONFIG_FILE="custom_iso.conf"

# 函数：安装依赖
install_dependencies() {
  echo "正在安装依赖..."
  local dependencies=("debootstrap" "squashfs-tools" "xorriso" "grub2-common")
  for dep in "${dependencies[@]}"; do
    if ! command -v "$dep" &> /dev/null; then
      echo "安装 $dep..."
      apt-get install -y "$dep"
    else
      echo "$dep 已安装"
    fi
  done
  echo "依赖安装完成"
}

# 函数：显示菜单
show_menu() {
  echo "请选择一个选项："
  echo "1) 从零开始构建ISO"
  echo "2) 从配置文件导入构建ISO"
  read -p "请输入选项编号: " option
  case $option in
    1) build_from_scratch ;;
    2) build_from_config ;;
    *) echo "无效选项，退出脚本" && exit 1 ;;
  esac
}

# 函数：从零开始构建
build_from_scratch() {
  install_dependencies
  set_work_dir
  set_distro_name
  set_debian_codename
  set_architecture
  set_desktop_env
  set_kernel_source
  set_installer
  import_packages
  build_system_dir
  install_kernel_and_packages
  chroot_system_dir
  install_installer  # 在打包squashfs前安装安装器
  pack_squashfs
  build_iso_dir
  build_iso
  save_config
}

# 函数：从配置文件导入构建
build_from_config() {
  install_dependencies
  read -p "请输入配置文件路径: " config_file
  if [ -f "$config_file" ]; then
    source "$config_file"
    build_system_dir
    install_kernel_and_packages
    chroot_system_dir
    install_installer  # 在打包squashfs前安装安装器
    pack_squashfs
    build_iso_dir
    build_iso
  else
    echo "配置文件不存在，退出脚本"
    exit 1
  fi
}

# 函数：设置工作目录
set_work_dir() {
  read -p "请输入工作目录路径: " work_dir
  mkdir -p "$work_dir"
  cd "$work_dir" || exit
  echo "工作目录设置为: $work_dir"
}

# 函数：设置发行版名称
set_distro_name() {
  read -p "请输入发行版名称: " distro_name
  echo "发行版名称设置为: $distro_name"
}

# 函数：选择Debian版本代号
set_debian_codename() {
  read -p "请输入Debian版本代号（如buster, bullseye）: " debian_codename
  echo "Debian版本代号设置为: $debian_codename"
}

# 函数：选择架构
set_architecture() {
  read -p "请输入架构（如amd64, i386）: " architecture
  echo "架构设置为: $architecture"
}

# 函数：选择桌面环境和登录管理器
set_desktop_env() {
  read -p "请输入桌面环境（如gnome, kde, xfce）: " desktop_env
  read -p "请输入登录管理器（如gdm3, sddm, lightdm）: " login_manager
  echo "桌面环境设置为: $desktop_env"
  echo "登录管理器设置为: $login_manager"
}

# 函数：选择内核来源
set_kernel_source() {
  read -p "是否使用用户提供的内核？(y/n): " use_custom_kernel
  if [ "$use_custom_kernel" = "y" ]; then
    read -p "请输入内核deb包目录路径: " kernel_dir
    echo "使用自定义内核目录: $kernel_dir"
  else
    echo "使用软件源内的内核"
  fi
}

# 函数：选择安装器
set_installer() {
  read -p "请选择安装器（如calamares, ubiquity）: " installer
  echo "安装器设置为: $installer"
}

# 函数：导入软件包
import_packages() {
  read -p "请输入要安装的软件包名称（用空格分隔）: " packages
  read -p "请输入包含deb包的文件夹路径: " deb_folder
  mkdir -p "$work_dir/packages"
  if [ -n "$deb_folder" ]; then
    cp -r "$deb_folder"/*.deb "$work_dir/packages/"
    echo "deb包已复制到: $work_dir/packages/"
  fi
  echo "软件包设置为: $packages"
}

# 函数：构建系统目录
build_system_dir() {
  echo "正在构建系统目录..."
  mkdir -p "$work_dir/system"
  # 使用debootstrap创建基本系统
  debootstrap --arch="$architecture" "$debian_codename" "$work_dir/system"
  echo "系统目录构建完成: $work_dir/system"
}

# 函数：安装内核和软件包
install_kernel_and_packages() {
  echo "正在安装内核和软件包..."
  # 挂载必要的文件系统
  mount --bind /proc "$work_dir/system/proc"
  mount --bind /sys "$work_dir/system/sys"
  mount --bind /dev "$work_dir/system/dev"
  
  # 配置APT源
  cat <<EOF > "$work_dir/system/etc/apt/sources.list"
deb http://deb.debian.org/debian $debian_codename main contrib non-free
deb http://deb.debian.org/debian $debian_codename-updates main contrib non-free
deb http://security.debian.org/debian-security $debian_codename-security main contrib non-free
EOF
  
  # 更新APT缓存
  chroot "$work_dir/system" apt-get update
  
  # 安装用户指定的软件包
  if [ -n "$packages" ]; then
    chroot "$work_dir/system" apt-get install -y $packages
  fi
  
  # 安装内核
  if [ "$use_custom_kernel" = "y" ] && [ -d "$kernel_dir" ]; then
    # 安装自定义内核
    cp "$kernel_dir"/*.deb "$work_dir/system/tmp/"
    chroot "$work_dir/system" dpkg -i /tmp/*.deb
  else
    # 安装Debian源内的最新内核
    chroot "$work_dir/system" apt-get install -y linux-image-$architecture linux-headers-$architecture
  fi
  
  # 安装桌面环境和登录管理器
  if [ -n "$desktop_env" ]; then
    chroot "$work_dir/system" apt-get install -y $desktop_env $login_manager
  fi
  
  # 清理APT缓存
  chroot "$work_dir/system" apt-get clean
  
  echo "内核和软件包安装完成"
}

# 函数：进入chroot环境
chroot_system_dir() {
  read -p "是否要chroot进入系统目录进行手动配置？(y/n): " chroot_confirm
  if [ "$chroot_confirm" = "y" ]; then
    echo "正在进入chroot环境..."
    chroot "$work_dir/system" /bin/bash
    echo "退出chroot环境"
  fi
  # 卸载挂载点
  umount "$work_dir/system/proc"
  umount "$work_dir/system/sys"
  umount "$work_dir/system/dev"
}

# 函数：安装安装器
install_installer() {
  if [ -n "$installer" ]; then
    echo "正在安装安装器: $installer..."
    # 挂载必要的文件系统
    mount --bind /proc "$work_dir/system/proc"
    mount --bind /sys "$work_dir/system/sys"
    mount --bind /dev "$work_dir/system/dev"
    
    # 安装安装器
    chroot "$work_dir/system" apt-get install -y "$installer"
    
    # 卸载挂载点
    umount "$work_dir/system/proc"
    umount "$work_dir/system/sys"
    umount "$work_dir/system/dev"
    
    echo "安装器 $installer 安装完成"
  fi
}

# 函数：打包系统目录为squashfs
pack_squashfs() {
  echo "正在打包系统目录为squashfs..."
  # 清理临时文件
  rm -rf "$work_dir/system/tmp/*"
  # 打包系统目录
  mksquashfs "$work_dir/system" "$work_dir/filesystem.squashfs" -comp xz
  # 计算文件大小
  du -sx --block-size=1 "$work_dir/system" | cut -f1 > "$work_dir/filesystem.size"
  echo "squashfs打包完成: $work_dir/filesystem.squashfs"
}

# 函数：构建ISO镜像目录
build_iso_dir() {
  echo "正在构建ISO镜像目录..."
  mkdir -p "$work_dir/iso"
  # 创建ISO目录结构
  mkdir -p "$work_dir/iso/casper"
  mkdir -p "$work_dir/iso/boot/grub"
  
  # 复制squashfs文件
  cp "$work_dir/filesystem.squashfs" "$work_dir/iso/casper/"
  cp "$work_dir/filesystem.size" "$work_dir/iso/casper/"
  
  # 检测并复制vmlinuz和initrd文件
  vmlinuz_path=$(find "$work_dir/system/boot" -name "vmlinuz-*" | head -n 1)
  initrd_path=$(find "$work_dir/system/boot" -name "initrd.img-*" | head -n 1)
  
  if [ -z "$vmlinuz_path" ] || [ -z "$initrd_path" ]; then
    echo "错误：未找到vmlinuz或initrd文件"
    exit 1
  fi
  
  cp "$vmlinuz_path" "$work_dir/iso/casper/vmlinuz"
  cp "$initrd_path" "$work_dir/iso/casper/initrd"
  
  # 创建GRUB配置文件
  cat <<EOF > "$work_dir/iso/boot/grub/grub.cfg"
set default="0"
set timeout=10
menuentry "$distro_name" {
  linux /casper/vmlinuz boot=casper quiet splash --
  initrd /casper/initrd
}
EOF
  echo "ISO镜像目录构建完成: $work_dir/iso"
}

# 函数：生成ISO镜像
build_iso() {
  read -p "请输入ISO文件名称: " iso_name
  echo "正在生成ISO镜像: $iso_name.iso..."
  # 使用xorriso生成ISO
  xorriso -as mkisofs \
    -r -V "$distro_name" \
    -o "$work_dir/$iso_name.iso" \
    -J -l -b boot/grub/grub.img \
    -no-emul-boot -boot-load-size 4 -boot-info-table \
    --grub2-boot-info --grub2-mbr /usr/lib/grub/i386-pc/boot_hybrid.img \
    "$work_dir/iso"
  if [ $? -eq 0 ]; then
    echo "ISO镜像生成成功: $work_dir/$iso_name.iso"
  else
    echo "ISO镜像生成失败"
    exit 1
  fi
}

# 函数：保存配置文件
save_config() {
  echo "保存配置文件到: $CONFIG_FILE"
  cat <<EOF > "$CONFIG_FILE"
work_dir="$work_dir"
distro_name="$distro_name"
debian_codename="$debian_codename"
architecture="$architecture"
desktop_env="$desktop_env"
login_manager="$login_manager"
use_custom_kernel="$use_custom_kernel"
kernel_dir="$kernel_dir"
installer="$installer"
packages="$packages"
deb_folder="$deb_folder"
EOF
}

# 主程序
show_menu